import torch
import config
import torch.nn as nn  # 导入PyTorch神经网络模块


class YoloV1Loss (nn.Module) :
    def __init__ (self, batch_size=config.BATCH_SIZE, S=config.S, B=config.B, C=config.C, l_coord=5, l_noobj=0.5) :
        """
        初始化YOLOv1损失函数
        :param batch_size: 批量大小，用于计算平均损失
        :param S: 网格数量（S×S），对应config中的S
        :param B: 每个网格的边界框数，对应config中的B
        :param C: 类别数，对应config中的C
        :param l_coord: 坐标损失权重（YOLOv1论文中设为5）
        :param l_noobj: 无目标置信度损失权重（YOLOv1论文中设为0.5）
        """
        super (YoloV1Loss, self).__init__ ()
        self.S = S  # 网格数量S
        self.B = B  # 每个网格的边界框数B
        self.C = C  # 类别数C
        self.depth = self.B * 5 + self.C  # 输出张量深度（2×5属性 + 20类别）
        self.batch_size = batch_size  # 批量大小
        self.l_coord = l_coord  # 坐标损失权重（平衡坐标误差和置信度误差）
        self.l_noobj = l_noobj  # 无目标置信度损失权重（降低无目标网格的损失贡献）

    def forward (self, preds, labels) :
        """
        计算预测值与标签的总损失（YOLOv1损失函数核心）
        :param preds: 模型预测输出（形状：(batch, S, S, depth)，默认(128,7,7,30)）
        :param labels: 真实标签（形状同preds）
        :return: 批量内每个样本的平均损失
        """
        coord_loss = self.cal_coord_loos (preds, labels)  # 计算坐标位置损失（中心坐标误差）
        bbox_loss = self.cal_bbox_loss (preds, labels)    # 计算边界框尺寸损失（宽高误差）
        obj_conf_loss = self.cal_obj_conf_loss (preds, labels)  # 计算有目标置信度损失
        noobj_conf_loss = self.cal_noobj_conf_loss (preds, labels)  # 计算无目标置信度损失
        class_loss = self.cal_class_loss (preds, labels)  # 计算类别分类损失
        # 总损失 = 坐标损失权重*(坐标位置损失+边界框尺寸损失) + 有目标置信度损失 + 无目标置信度损失权重*无目标置信度损失 + 类别损失
        total_loss = self.l_coord * coord_loss + self.l_coord * bbox_loss + \
                     obj_conf_loss + self.l_noobj * noobj_conf_loss + class_loss
        loss_of_each_data = torch.sum (total_loss) / self.batch_size  # 计算批量内平均损失
        return loss_of_each_data

    def cal_coord_loos (self, pred, label) :
        """
        计算边界框中心坐标的平方误差（对应YOLOv1损失函数中的坐标位置损失）
        :param pred: 模型预测张量（形状：(batch, S, S, depth)）
        :param label: 真实标签张量（形状同pred）
        :return: 坐标位置损失（形状：(batch_size, 1)）
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')  # 初始化损失张量（设备与输入一致）

        ## 说明：仅对存在目标的网格计算坐标损失（label[:, :, :, 0]为1表示该网格有目标）
        # 向量运算计算第一个边界框的中心坐标误差（x,y）
        loss += torch.sum (
            label[:, :, :, 0] *  # 掩码（仅保留有目标的网格）
            (
                    (label[:, :, :, 1] - pred[:, :, :, 1]) ** 2 +  # x坐标平方差
                    (label[:, :, :, 2] - pred[:, :, :, 2]) ** 2    # y坐标平方差
            )
        ) + torch.sum (
            label[:, :, :, 5] *  # 第二个边界框的存在掩码（YOLOv1中每个网格预测2个边界框）
            (
                    (label[:, :, :, 6] - pred[:, :, :, 6]) ** 2 +  # 第二个边界框x坐标平方差
                    (label[:, :, :, 7] - pred[:, :, :, 7]) ** 2    # 第二个边界框y坐标平方差
            )
        )

        return loss  # 返回坐标位置损失

    def cal_bbox_loss (self, pred, label) :
        """
        计算边界框宽高的平方误差（对应YOLOv1损失函数中的边界框尺寸损失）
        :param pred: 模型预测张量
        :param label: 真实标签张量
        :return: 边界框尺寸损失（形状：(batch_size, 1)）
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')  # 初始化损失张量

        # 计算第一个边界框的宽高误差（w,h）
        loss += torch.sum (  # 第一个边界框
            label[:, :, :, 0] *  # 有目标网格掩码
            (
                    (label[:, :, :, 3] - pred[:, :, :, 3]) ** 2 +  # 宽度平方差
                    (label[:, :, :, 4] - pred[:, :, :, 4]) ** 2     # 高度平方差
            )
        ) + torch.sum (  # 第二个边界框
            label[:, :, :, 5] *  # 第二个边界框存在掩码
            (
                    (label[:, :, :, 8] - pred[:, :, :, 8]) ** 2 +  # 第二个边界框宽度平方差
                    (label[:, :, :, 9] - pred[:, :, :, 9]) ** 2     # 第二个边界框高度平方差
            )
        )

        return loss  # 返回边界框尺寸损失

    def cal_obj_conf_loss (self, pred, label) :
        """
        计算有目标网格的置信度损失（预测置信度与真实置信度的平方差）
        :param pred: 模型预测张量（置信度在第0和第5通道）
        :param label: 真实标签张量（置信度为1表示有目标）
        :return: 有目标置信度损失（形状：(batch_size, 1)）
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')  # 初始化损失张量

        # 有目标网格的置信度损失：预测置信度与1（存在目标）的平方差
        loss += torch.sum (
            label[:, :, :, 0] * (label[:, :, :, 0] - pred[:, :, :, 0]) ** 2  # 第一个边界框置信度损失
        )

        return loss  # 返回有目标置信度损失

    def cal_noobj_conf_loss (self, pred, label) :
        """
        计算无目标网格的置信度损失（降低无目标网格的置信度预测值）
        :param pred: 模型预测张量
        :param label: 真实标签张量（置信度为0表示无目标）
        :return: 无目标置信度损失（形状：(batch_size, 1)）
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')  # 初始化损失张量

        # 无目标网格掩码：通过-(label[:, :, :, 0] - 1)将有目标网格（1）变为0，无目标网格（0）变为1
        loss += torch.sum (
            (-(label[:, :, :, 0] - 1)) * (label[:, :, :, 0] - pred[:, :, :, 0]) ** 2
        )
        # 说明：无目标网格的真实置信度为0，计算预测置信度与0的平方差

        return loss  # 返回无目标置信度损失

    def cal_class_loss (self, pred, label) :
        """
        计算类别分类损失（仅对有目标的网格计算类别概率的平方差）
        :param pred: 模型预测张量（类别概率在第10-29通道）
        :param label: 真实标签张量（类别为独热编码）
        :return: 类别分类损失（形状：(batch_size, 1)）
        """
        loss = torch.zeros ((self.batch_size, 1)).to ('cuda')  # 初始化损失张量

        # 有目标网格的类别损失：预测类别概率与真实独热编码的平方差之和
        loss += torch.sum (
            label[:, :, :, 0] * (
                torch.sum ((label[:, :, :, 10 :30] - pred[:, :, :, 10 :30]) ** 2, dim=3)  # 沿类别维度求和平方差
            )
        )
        # 说明：仅当网格有目标时（label[:, :, :, 0]=1），才计算类别损失

        return loss  # 返回类别分类损失


if __name__ == '__main__' :
    pred = torch.zeros ((2, 3, 3, 5))
    label = torch.randn ((2, 3, 3, 5))
    loss = torch.zeros ((2, 1))
    print (loss.shape)
    print (torch.sum (pred).shape)
    loss += torch.sum (pred[:, :, :, 0] - label[:, :, :, 0])
    print (loss)
